import numpy as np
import matplotlib.pyplot as plt

from hartreefock import *

def do_HeH_calculation():
    '''
    HeH+ 在 R = 1.4632 时的 Hartree-Fock 计算
    '''
    mol = HeH(1.4632)

    print('start HeH+ Hartree-Fock claculation, with convergence criterion: δE = 1e-10 a.u.')
    E0, ei, basis, Cx = HartreeFock(mol, details=True, verbose=1)

    print('calculation result of HeH+ at a sparation of R = 1.4632 a.u.')
    print(f'ground state energy = {E0:.3f} a.u., or {E0*hartree:.3f} eV')
    print(f'orbital energies:')
    for i, e in enumerate(ei):
        print(f'   {i}: {e:.5f} a.u.')

    z = np.linspace(-1.5, 2.5)
    r = np.linspace(-2, 2)
    z, r = np.meshgrid(z, r)
    d = eval_cloud(basis, Cx[:, 0], z, r)

    plt.figure()
    ax = plt.axes(projection='3d')
    ax.plot_surface(z, r, d)
    plt.xlabel(r'$z/\mathrm{a.u.}$')
    plt.ylabel(r'$r/\mathrm{a.u.}$')
    plt.savefig("HeH.png")


def do_H2_calculation():
    '''
    计算 H2 的能量曲线
    '''
    r = np.linspace(1, 2)
    E = []
    for ri in r:
        E.append(HartreeFock(H2(ri)))
    E = np.array(E)

    plt.figure()
    plt.plot(r, E)
    plt.xlabel(r'$r/\mathrm{a.u.}$')
    plt.ylabel(r'$E/\mathrm{a.u.}$')
    plt.savefig("H2.png")


if __name__ == '__main__':
    do_HeH_calculation()
    do_H2_calculation()
